/*
 * PopTree.h
 *
 *  Created on: May 28, 2012
 *      Author: pickrell
 */

#ifndef POPTREE_H_
#define POPTREE_H_


#include "PhyloPop_params.h"
#include "Alignment.h"
using namespace std;

struct Node
{
	int index;
	string name;
	double height;
	bool is_tip;
	bool is_root;
	bool rev;
	int pop;
	float lm;
	vector<float> l1;
	vector<float> l0;
};

struct Dist
{
	float weight;
	float len;
};

typedef adjacency_list<listS, listS, boost::bidirectionalS, Node, Dist> Graph;
typedef pair<int, int> Edge;
typedef boost::property_map<Graph, int Node::*>::type IndexMap;
typedef boost::graph_traits<Graph>::vertex_iterator vertex_iter;

class PopTree{
public:

	// initialize the graph
	PopTree();

	int nsite, len, nmono, ntip;
	vector<int> colcounts;
	double phi, resphi;
	// initialize from a Newick string
	PopTree(string);
	void set_graph(string); //or set it to a Newick string
	void set_heights_from_edges();

	// initialize with three populations
	PopTree(vector<string>, vector<vector<bool> >, int, int, vector<int>, PhyloPop_params *);


	PhyloPop_params * params;
	//copy from another location
	void copy(PopTree*);
	//void copy_helper(Graph::vertex_descriptor, Graph::vertex_descriptor, PopGraph *);

	Graph g;
	int indexcounter;

	//get the depth of the node
	float get_coaltime(Graph::vertex_descriptor);

	vector<Graph::vertex_descriptor> index2father;
	vector<string> popnames;
	Graph::vertex_descriptor root;
	map<int, Graph::vertex_descriptor> index2vertex();

	Graph::vertex_descriptor add_tip(Graph::vertex_descriptor, vector<bool>, string);
	void remove_tip(Graph::vertex_descriptor);


	void set_root(Graph::vertex_descriptor);
	set<Graph::vertex_descriptor> get_root_adj();
	set<Graph::edge_descriptor> get_root_adj_edge();
	//void print();
	void print(string);

	map<string, Graph::vertex_descriptor> get_tips( Graph::vertex_descriptor);

	double get_height(Graph::vertex_descriptor);
	double get_dist_to_root(Graph::vertex_descriptor); //get the distance to the root for any node in the tree
	void flip_sons(Graph::vertex_descriptor, gsl_rng* ); //flip the inorder traversal from a node

	//get the inorder traversal
	vector< Graph::vertex_descriptor > get_inorder_traversal( int);
	void inorder_traverse(Graph::vertex_descriptor, int*, vector<Graph::vertex_descriptor >*);
	vector< Graph::vertex_descriptor > get_inorder_traversal_noroot( int);
	vector< Graph::vertex_descriptor > get_inorder_traversal_fromindex( int);

	// postorder traversal
	vector< Graph::vertex_descriptor > get_postorder_traversal( int);
	vector<Graph::vertex_descriptor> get_postorder_traversal(Graph::vertex_descriptor, int);
	void postorder_traverse(Graph::vertex_descriptor, int*, vector<Graph::vertex_descriptor >*);

	//likelihoods
	double current_llk;
	double conditional_llk(Graph::vertex_descriptor);
	void optimize_height(Graph::vertex_descriptor);
	void optimize_height_gs(Graph::vertex_descriptor);
	void optimize_heights(vector<Graph::vertex_descriptor>);
	void optimize_all_heights();
	double quad_opt_h(double, double, double, double, double, double, double, double);
	int golden_section_h(Graph::vertex_descriptor, Graph::vertex_descriptor, double, double, double, double, int*);
	void compute_tree_llk();
	void compute_tree_llk(Graph::vertex_descriptor);
	double lk(int, int, double);
	void set_ls(Graph::vertex_descriptor);
	double sumlog(double, double);
	//get the LCA of two tips in the tree
	Graph::vertex_descriptor get_LCA(Graph::vertex_descriptor, Graph::vertex_descriptor, Graph::vertex_descriptor);

	//get the path to the root for a vertex
	set<Graph::vertex_descriptor> get_path_to_root(Graph::vertex_descriptor);
	vector<Graph::vertex_descriptor> get_path_to_root_vec(Graph::vertex_descriptor);

	//local rearrangement
	void local_rearrange(Graph::vertex_descriptor, int);
	void move_root(int);
	void move_root(gsl_rng*);

	void set_UPGMA(Alignment*);
	double set_random_UPGMA(Alignment*, double, double, double, gsl_rng *);
	pair<pair<int, int>, double> get_mindist(vector<Graph::vertex_descriptor>* , vector<vector<double> >*);
	pair<pair<int, int>, double> get_randdist(vector<Graph::vertex_descriptor>* , vector<vector<double> >*, double, gsl_rng *, double*);
	double get_postd(double, double, int, double, double, gsl_rng*, double *);
	double get_loggamma(double, double, double);

	//global rearrangement
	void global_rearrange(Graph::vertex_descriptor, Graph::vertex_descriptor);

	//Newick format
	string get_newick_format();
	string get_newick_format(Graph::vertex_descriptor);
	string get_newick_format(map<string, double>*);
	//string get_newick_subtrees(Graph::vertex_descriptor, Graph::vertex_descriptor); //join two subtrees (for migrations)
	void newick_helper( Graph::vertex_descriptor, string*);
	void newick_helper( Graph::vertex_descriptor, string*, int);

	pair<Graph::vertex_descriptor, Graph::vertex_descriptor> get_child_nodes(Graph::vertex_descriptor); // skips over migration nodes

	pair<Graph::vertex_descriptor, double> get_parent_node(Graph::vertex_descriptor); //skips over migration nodes
	void set_edge_lens();


};


#endif /* POPTREE_H_ */
